from flask import Flask, render_template, request, jsonify
import os
import requests
from concurrent.futures import ThreadPoolExecutor

app = Flask(__name__)

# 确保下载目录存在
DOWNLOAD_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'downloaded_images')
os.makedirs(DOWNLOAD_DIR, exist_ok=True)

def download_image(url, index):
    try:
        response = requests.get(url, stream=True)
        if response.status_code == 200:
            file_extension = os.path.splitext(url)[-1] or '.jpg'
            filename = f'image_{index + 1}{file_extension}'
            filepath = os.path.join(DOWNLOAD_DIR, filename)
            
            with open(filepath, 'wb') as f:
                for chunk in response.iter_content(chunk_size=8192):
                    if chunk:
                        f.write(chunk)
            return {'success': True, 'filename': filename, 'message': f'Successfully downloaded {filename}'}
        return {'success': False, 'message': f'Failed to download image {index + 1}: Status code {response.status_code}'}
    except Exception as e:
        return {'success': False, 'message': f'Error downloading image {index + 1}: {str(e)}'}

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/download', methods=['POST'])
def download_images():
    data = request.json
    if not data or 'urls' not in data or not data['urls']:
        return jsonify({'success': False, 'message': 'No URLs provided'})

    urls = data['urls']
    results = []
    
    # 使用线程池并行下载图片
    with ThreadPoolExecutor(max_workers=4) as executor:
        futures = [executor.submit(download_image, url, i) for i, url in enumerate(urls)]
        for future in futures:
            results.append(future.result())

    # 统计成功和失败的数量
    success_count = sum(1 for r in results if r['success'])
    total_count = len(urls)
    
    return jsonify({
        'success': True,
        'message': f'Downloaded {success_count} of {total_count} images',
        'results': results
    })

if __name__ == '__main__':
    app.run(debug=True)
